import os
import numpy as np
from pathlib import Path
from typing import List, Tuple
from PIL import Image
import imageio.v2 as imageio


def collect_pairs(out_dir: Path) -> List[Path]:
    cases = [p for p in out_dir.glob('*_cases.png')]
    code = [p for p in out_dir.glob('*_code.png')]
    index = {}
    for p in cases:
        key = p.name.replace('_cases.png', '')
        index.setdefault(key, {})['cases'] = p
    for p in code:
        key = p.name.replace('_code.png', '')
        index.setdefault(key, {})['code'] = p

    def sort_key(key: str) -> Tuple[int, float]:
        # key example: RID_PID_STAMP
        parts = key.split('_')
        stamp = parts[-1] if parts else ''
        ts = 0
        if stamp.isdigit():
            try:
                ts = int(stamp)
            except Exception:
                ts = 0
        # fallback to earliest file mtime among cases/code
        mt = 0.0
        v = index.get(key, {})
        files = [v.get('cases'), v.get('code')]
        mt_candidates = [f.stat().st_mtime for f in files if f]
        if mt_candidates:
            mt = min(mt_candidates)
        return (ts, mt)

    ordered_keys = sorted(index.keys(), key=sort_key)
    pairs: List[Path] = []
    for key in ordered_keys:
        v = index[key]
        if 'cases' in v:
            pairs.append(v['cases'])
        if 'code' in v:
            pairs.append(v['code'])
    return pairs


def pad_to_3840x2160(img: Image.Image) -> np.ndarray:
    # 等比缩放到宽不超过3840、高不超过2160，然后居中填充到3840x2160（4K）
    w, h = img.size
    scale = min(3840 / w, 2160 / h)
    nw = int(w * scale)
    nh = int(h * scale)
    img = img.resize((nw, nh), Image.LANCZOS)
    canvas = Image.new('RGB', (3840, 2160), (255, 255, 255))
    ox = (3840 - nw) // 2
    oy = (2160 - nh) // 2
    canvas.paste(img, (ox, oy))
    return np.array(canvas)


def main():
    root = Path.cwd()
    out_dir = Path(os.environ.get('OUTPUT_DIR', str(root / 'artifacts' / 'luogu-322081' / 'images')))
    out_dir.mkdir(parents=True, exist_ok=True)
    images = collect_pairs(out_dir)
    if not images:
        print('No images found for video')
        return
    out_path = out_dir.parent / 'luogu_records.mp4'
    writer = imageio.get_writer(str(out_path), fps=24, codec='libx264')
    try:
        fpi = int(os.environ.get('FRAMES_PER_IMAGE', '2'))
        for p in images:
            img = Image.open(p).convert('RGB')
            frame = pad_to_3840x2160(img)
            for _ in range(max(1, fpi)):
                writer.append_data(frame)
        print('Video written to', out_path)
    finally:
        writer.close()


if __name__ == '__main__':
    main()


if __name__ == '__main__':
    main()
